#include "initial_alignment.hpp"

void solveGyrBais(std::map<double, ImageFrame> &all_image_frames, Eigen::Vector3d* Bgs){
    Eigen::Matrix3d A;
    Eigen::Vector3d b;
    A.setZero(); b.setZero();
    std::map<double, ImageFrame>::iterator frame_i, frame_j;
    for(frame_i = all_image_frames.begin(); std::next(frame_i) != all_image_frames.end(); ++frame_i){
        frame_j = std::next(frame_i);
        Eigen::Matrix3d temp_A;
        Eigen::Vector3d temp_b;
        temp_A.setZero(); temp_b.setZero();
        temp_A = frame_j->second.pre_integration->jacobian.block<3, 3>(3, 12);
        temp_b = 2 * (frame_j->second.pre_integration->deltaQ.inverse() *
                      Eigen::Quaterniond(frame_i->second.R.transpose() * frame_j->second.R)).vec();
//            LOG(INFO) << "temp_A: " << temp_A << std::endl;
//            LOG(INFO) << "temp_b: " << temp_b << std::endl;
        A += temp_A.transpose() * temp_A;
        b += temp_A.transpose() * temp_b;
    }
//        LOG(INFO) << "A: " << A << std::endl;
//        LOG(INFO) << "b: " << b << std::endl;
    Eigen::Vector3d dx = A.ldlt().solve(b);
    LOG(INFO) << "dx: " << dx.transpose() << std::endl;

    for(int i = 0; i <= WINDOW_SIZE; i++){
        Bgs[i] += dx;
    }

    for(frame_i = all_image_frames.begin(); std::next(frame_i) != all_image_frames.end(); frame_i++){
        frame_j = std::next(frame_i);
        frame_j->second.pre_integration->repropagate(Eigen::Vector3d::Zero(), Bgs[0]);
    }

}

void secondAlignment(std::map<double, ImageFrame> &all_image_frames, Eigen::Vector3d& g, Eigen::VectorXd& x){

    int n = all_image_frames.size();
    int N = 3 * n + 2 + 1;
    Eigen::MatrixXd A(N, N);
    Eigen::VectorXd b(N);
    A.setZero(); b.setZero();
    Eigen::Vector3d g0 = G.norm() * g.normalized();
    for(int k = 0; k < 1; k++){
        Eigen::Vector3d w1, w2;
        Eigen::Matrix<double, 3, 2> w1_w2;
        {
            Eigen::Vector3d w0 = g0.normalized();
            Eigen::Vector3d temp;
            temp << 0, 0, 1;
            if(w0 == temp){
                temp << 1, 0, 0;
            }
            //施密特正交化
            w1 = (temp - (w0.transpose() * temp) * w0).normalized();
            w2 = w0.cross(w1);
            w1_w2.block<3, 1>(0, 0) = w1;
            w1_w2.block<3, 1>(0, 1) = w2;
        }

        std::map<double, ImageFrame>::iterator frame_i, frame_j;
        int i = 0;
        for(frame_i = all_image_frames.begin(); std::next(frame_i) != all_image_frames.end(); ++frame_i, ++i){
            frame_j = std::next(frame_i);
            Eigen::MatrixXd temp_A(6, 9);
            Eigen::VectorXd temp_b(6);
            double dt = frame_j->second.pre_integration->sum_dt;
            temp_A.setZero(); temp_b.setZero();
            temp_A.block<3, 3>(0, 0) = -Eigen::Matrix3d::Identity() * dt;
            temp_A.block<3, 2>(0, 6) = 0.5 * frame_i->second.R.transpose() * dt * dt * w1_w2;
            temp_A.block<3, 1>(0, 8) = frame_i->second.R.transpose() * (frame_j->second.T - frame_i->second.T) / 100.0;
            temp_b.block<3, 1>(0, 0) = frame_j->second.pre_integration->deltaP + frame_i->second.R.transpose() * frame_j->second.R * TIC[0] - TIC[0]
                                       - 0.5 * frame_i->second.R.transpose() * dt * dt * g0;

            temp_A.block<3, 3>(3, 0) = -Eigen::Matrix3d::Identity();
            temp_A.block<3, 3>(3, 3) = frame_i->second.R.transpose() * frame_j->second.R;
            temp_A.block<3, 2>(3, 6) = frame_i->second.R.transpose() * dt * w1_w2;
            temp_b.block<3, 1>(3, 0) = frame_j->second.pre_integration->deltaV - frame_i->second.R.transpose() * dt * g0;

            Eigen::Matrix<double, 6, 6> cov_inv = Eigen::Matrix<double, 6, 6>::Zero();
            //cov.block<6, 6>(0, 0) = IMU_cov[i + 1];
            //MatrixXd cov_inv = cov.inverse();
            cov_inv.setIdentity();

            Eigen::MatrixXd r_A = temp_A.transpose() * cov_inv * temp_A;
            Eigen::VectorXd r_b = temp_A.transpose() * cov_inv * temp_b;

            A.block<6, 6>(i * 3, i * 3) += r_A.topLeftCorner<6, 6>();
            b.block<6, 1>(i * 3, 0) += r_b.block<6, 1>(0, 0);

            A.bottomRightCorner<3, 3>() += r_A.bottomRightCorner<3, 3>();
            b.block<3, 1>(N - 3, 0) += r_b.block<3, 1>(6, 0);

            A.block<6, 3>(i * 3, N - 3) += r_A.topRightCorner<6, 3>();
            A.block<3, 6>(N - 3, i * 3) += r_A.bottomLeftCorner<3, 6>();
        }
        A = A * 1000.0;
        b = b * 1000.0;
        x = A.ldlt().solve(b);
        Eigen::Vector2d dg = x.block<2, 1>(N - 3, 0);
        Eigen::Vector3d g_temp = G.norm() * (g0 + w1_w2 * dg).normalized();
        g0 = g_temp;
    }
    g = g0;

}

bool preAlignment(std::map<double, ImageFrame> &all_image_frames, Eigen::Vector3d& g, Eigen::VectorXd& x){
    int n = all_image_frames.size();
    int N = 3 * n + 4;
    Eigen::MatrixXd A(N, N);
    Eigen::VectorXd b(N);
    A.setZero(); b.setZero();
    std::map<double, ImageFrame>::iterator frame_i, frame_j;
    int i = 0;
    for(frame_i = all_image_frames.begin(); std::next(frame_i) != all_image_frames.end(); ++frame_i, ++i){
        frame_j = std::next(frame_i);
        Eigen::MatrixXd temp_A(6, 10);
        Eigen::VectorXd temp_b(6);
        double dt = frame_j->second.pre_integration->sum_dt;
        temp_A.setZero(); temp_b.setZero();
        temp_A.block<3, 3>(0, 0) = -Eigen::Matrix3d::Identity() * dt;
        temp_A.block<3, 3>(0, 6) = 0.5 * frame_i->second.R.transpose() * dt * dt;
        temp_A.block<3, 1>(0, 9) = frame_i->second.R.transpose() * (frame_j->second.T - frame_i->second.T) / 100.0;
        temp_b.block<3, 1>(0, 0) = frame_j->second.pre_integration->deltaP + frame_i->second.R.transpose() * frame_j->second.R * TIC[0] - TIC[0];

        temp_A.block<3, 3>(3, 0) = -Eigen::Matrix3d::Identity();
        temp_A.block<3, 3>(3, 3) = frame_i->second.R.transpose() * frame_j->second.R;
        temp_A.block<3, 3>(3, 6) = frame_i->second.R.transpose() * dt;
        temp_b.block<3, 1>(3, 0) = frame_j->second.pre_integration->deltaV;

        Eigen::Matrix<double, 6, 6> cov_inv = Eigen::Matrix<double, 6, 6>::Zero();
        //cov.block<6, 6>(0, 0) = IMU_cov[i + 1];
        //MatrixXd cov_inv = cov.inverse();
        cov_inv.setIdentity();

        Eigen::MatrixXd r_A = temp_A.transpose() * cov_inv * temp_A;
        Eigen::VectorXd r_b = temp_A.transpose() * cov_inv * temp_b;

        A.block<6, 6>(i * 3, i * 3) += r_A.topLeftCorner<6, 6>();
        b.block<6, 1>(i * 3, 0) += r_b.block<6, 1>(0, 0);

        A.bottomRightCorner<4, 4>() += r_A.bottomRightCorner<4, 4>();
        b.block<4, 1>(N - 4, 0) += r_b.block<4, 1>(6, 0);

        A.block<6, 4>(i * 3, N - 4) += r_A.topRightCorner<6, 4>();
        A.block<4, 6>(N - 4, i * 3) += r_A.bottomLeftCorner<4, 6>();
    }
    A = A * 1000.0;
    b = b * 1000.0;
    x = A.ldlt().solve(b);
    double s = x(N - 1) / 100.0;
    LOG(INFO) << "S: " << s << std::endl;
    g  = x.block<3, 1>(N - 4, 0);
    LOG(INFO) << "g.norm: " << g.norm() << " g.transpose(): " << g.transpose() << std::endl;
    if(std::abs(g.norm() - G.norm()) > 1.0 || s < 0){
        return false;
    }
    secondAlignment(all_image_frames, g, x);
    s = (x.tail<1>())(0) / 100.0;
    (x.tail<1>())(0) = s;
    if(s < 0.0){
        return false;
    }
    else{
        return true;
    }
}
